%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                                                 %
%Replication Code for Baird Bohren McIntosh Ozler %
%                                                 %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

diary 'Table1.txt'
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                Table 1, Column 1                %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Clustered design
% Half of clusters pure control, half pure treatment, ICC = 0.0

clear; clc;
n     = 10;         % cluster size
C     = 100;        % number of clusters
tau   = 0;          % variance of cluster error (tau^2)
sigma = 1;          % variance of individual error (sigma^2)
pi    = [0,1];      % cluster saturations
f     = [1/2,1/2];  % fraction of clusters in each saturation, respectively

[SE_T, SE_S, SE_Tonly] = power_pooled(n,C,tau,sigma,pi,f);
SE_Tonly

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                Table 1, Column 2                %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Clustered design
% Half of clusters pure control, half pure treatment, ICC = 0.1

clear; clc;
n     = 10;        % cluster size
C     = 100;       % number of clusters
tau   = 0.1;       % variance of cluster error (tau^2)
sigma = 0.9;       % variance of individual error (sigma^2)
pi    = [0,1];     % cluster saturations
f     = [1/2,1/2]; % fraction of clusters in each saturation, respectively

[SE_T, SE_S, SE_Tonly] = power_pooled(n,C,tau,sigma,pi,f);
SE_Tonly

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                Table 1, Column 3                %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Clustered design
% Half of clusters pure control, half pure treatment, ICC = 1.0

clear; clc;
n     = 10;        % cluster size
C     = 100;       % number of clusters
tau   = 1;         % variance of cluster error    (tau^2)
sigma = 0;         % variance of individual error (sigma^2)
pi    = [0,1];     % cluster saturations
f     = [1/2,1/2]; % fraction of clusters in each saturation, respectively

[SE_T, SE_S, SE_Tonly] = power_pooled(n,C,tau,sigma,pi,f);

SE_Tonly

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                Table 1, Column 4                %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Partial Population design
% Treatment cluster has saturation 0.5, ICC = 0.0

clear; clc;
n     = 10;        % cluster size
C     = 100;       % number of clusters
tau   = 0;         % variance of cluster error (tau^2)
sigma = 1;         % variance of individual error (sigma^2)
pi    = [0,0.5];   % cluster saturations
psi   = [0.01:.001:.99];
SE_ITT = zeros(1,length(psi));
SE_SNT = SE_ITT;

for j = 1:length(psi);
    f = [psi(j),1-psi(j)]; 
    [SE_ITT(j), SE_SNT(j),z] = power_pooled(n,C,tau,sigma,pi,f);
end;

obj   = SE_ITT + SE_SNT;  % objective function to maximize
[i,j] = index_min(obj); % index of psi that minimizes objective

psi(i,j)                % optimal proportion of pure control clusters
1 - psi(i,j)            % optimal proportion of partial treatment clusters
SE_ITT(i,j)              % SE_ITT at optimal control size
SE_SNT(i,j)              % SE_SNT at optimal control size

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                Table 1, Column 5                %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Partial Population design
% Treatment cluster has saturation 0.5, ICC = 0.1

clear; clc;
n     = 10;        % cluster size
C     = 100;       % number of clusters
tau   = 0.1;       % variance of cluster error (tau^2)
sigma = 0.9;       % variance of individual error (sigma^2)
pi    = [0,0.5];   % cluster saturations
psi   = [0.01:.001:.99];
SE_ITT = zeros(1,length(psi));
SE_SNT = SE_ITT;

for j = 1:length(psi);%j=500
    f = [psi(j),1-psi(j)]; 
    [SE_ITT(j), SE_SNT(j),z] = power_pooled(n,C,tau,sigma,pi,f);
end;

obj   = SE_ITT + SE_SNT;  % objective function to maximize
[i,j] = index_min(obj); % index of psi that minimizes objective

psi(i,j)                % optimal proportion of pure control clusters
1 - psi(i,j)            % optimal proportion of partial treatment clusters
SE_ITT(i,j)              % SE_ITT at optimal control size
SE_SNT(i,j)              % SE_SNT at optimal control size

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                Table 1, Column 6                %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Partial Population design
% Treatment cluster has saturation 0.5, ICC = 1.0

clear; clc;
n     = 10;        % cluster size
C     = 100;       % number of clusters
tau   = 1;         % variance of cluster error (tau^2)
sigma = 0;         % variance of individual error (sigma^2)
pi    = [0,0.5];   % cluster saturations
psi   = [0.01:.001:.99];
SE_ITT = zeros(1,length(psi)); 
SE_SNT = SE_ITT;

for j = 1:length(psi);
    f = [psi(j),1-psi(j)]; 
    [SE_ITT(j), SE_SNT(j),z] = power_pooled(n,C,tau,sigma,pi,f);
end;

obj   = SE_ITT+SE_SNT;    % objective function to maximize
[i,j] = index_min(obj); % index of psi that minimizes objective

psi(i,j)                % optimal proportion of pure control clusters
1-psi(i,j)              % optimal proportion of partial treatment clusters
SE_ITT(i,j)              % SE_ITT at optimal control size
SE_SNT(i,j)              % SE_SNT at optimal control size % Begin Table 1, Column 6

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                Table 1, Column 7-8              %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Partial Population design
% Some clusters pure control, the rest partially treated, ICC = 0.1

clear; clc;
n     = 10;        % cluster size
C     = 100;       % number of clusters
tau   = 0.1;       % variance of cluster error (tau^2)
sigma = 0.9;       % variance of individual error (sigma^2)
x     = [0.01:.005:.99]; % values for optimization
SE_ITT = zeros(length(x),length(x)); 
SE_SNT = SE_ITT;

for i = 1:length(x);
    pi = [0,x(i)];
    for j = 1:length(x);
        f = [x(j),1-x(j)]; 
        [SE_ITT(i,j), SE_SNT(i,j),z] = power_pooled(n,C,tau,sigma,pi,f);
    end;
end;

%%%% Table 1, Column 7
% Subject to twice as much weight on SE_SNT compared to SE_ITT
obj   = SE_ITT + 2*SE_SNT;% objective function to maximize
[i,j] = index_min(obj); % index of psi that minimizes SE_ITT

x(i)                    % optimal saturation of treatment
x(j)                    % optimal proportion of pure control clusters
1-x(j)                  % optimal proportion of partial treatment clusters
SE_ITT(i,j)              % SE_ITT at optimal control size
SE_SNT(i,j)              % SE_SNT at optimal control size

%%%% Table 2, Column 8
% Subject to SE_ITT <= 0.09 (MDE_T <= .25)
SE_SNT_aug = SE_SNT;
for i = 1:length(x);
    for j = 1:length(x);
        if SE_ITT(i,j) > 0.09;
            SE_SNT_aug(i,j) = Inf;
        end;
    end;
end;
obj = SE_SNT_aug;        % objective function to maximize
[i,j] = index_min(obj); % index of psi that minimizes SE_SNT

x(i)                    % optimal saturation of treatment
x(j)                    % optimal proportion of pure control clusters
1 - x(j)                % optimal proportion of partial treatment clusters
SE_ITT(i,j)              % SE_ITT at optimal control size
SE_SNT(i,j)              % SE_SNT at optimal control size

diary off
diary 'Table2_13.txt'
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                Table 2, Column 1                %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Random saturation design (two partially treated clusters), ICC = 0.0
% Optimal cluster saturations and allocation

clear; clc;
n     = 10;        % cluster size
C     = 100;       % number of clusters
tau    = 0.0;      % variance of cluster error (tau^2)
sigma  = 1.0;      % variance of individual error (sigma^2)
x      = [0.01:.01:.5]; % values for optimization
SE_DT = zeros(length(x),length(x));
SE_DNT = SE_DT;

for i = 1:length(x);
    pi = [x(i),1-x(i)];
    for j = 1:length(x);
        f = [x(j),1-x(j)];
        [SE_DT(i,j),SE_DNT(i,j)] = power_slope(n,C,tau,sigma,pi,f,1,2);
    end;
end;

obj   = SE_DT + SE_DNT; % objective function to maximize
[i,j] = index_min(obj); % index of psi that minimizes objective
x(i)         % optimal saturation for lower saturation
1 - x(i)     % optimal saturation for higher saturation
x(j)         % optimal fraction of clusters in lower saturation
1 - x(j)     % optimal fraction of clusters in higher saturation
SE_DT(i,j)  % MDSE_T at optimal 
SE_DNT(i,j)  % MDSE_S at optimal

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                Table 2, Column 2                %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Random saturation design (two partially treated clusters), ICC = 0.1
% Optimal cluster saturations and allocation

clear; clc;
n     = 10;        % cluster size
C     = 100;       % number of clusters
tau    = 0.1;      % variance of cluster error (tau^2)
sigma  = 0.9;      % variance of individual error (sigma^2)
x      = [0.01:.01:.5]; % values for optimization
SE_DT = zeros(length(x),length(x));
SE_DNT = SE_DT;

for i = 1:length(x);
    pi = [x(i),1-x(i)];
    for j = 1:length(x);
        f = [x(j),1-x(j)];
        [SE_DT(i,j),SE_DNT(i,j)] = power_slope(n,C,tau,sigma,pi,f,1,2);
    end;
end;

obj   = SE_DT + SE_DNT; % objective function to maximize
[i,j] = index_min(obj); % index of psi that minimizes objective
x(i)         % optimal saturation for lower saturation
1 - x(i)     % optimal saturation for higher saturation
x(j)         % optimal fraction of clusters in lower saturation
1 - x(j)     % optimal fraction of clusters in higher saturation
SE_DT(i,j)  % MDSE_T at optimal 
SE_DNT(i,j)  % MDSE_S at optimal

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                Table 2, Column 3                %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Random saturation design (two partially treated clusters), ICC = 1.0
% Optimal cluster saturations and allocation

clear; clc;
n     = 10;        % cluster size
C     = 100;       % number of clusters
tau    = 1.0;      % variance of cluster error (tau^2)
sigma  = 0.0;      % variance of individual error (sigma^2)
x      = [0.1:.01:.5]; % values for optimization
SE_DT = zeros(length(x),length(x));
SE_DNT = SE_DT;

for i = 1:length(x);
    pi = [x(i),1-x(i)];
    for j = 1:length(x);
        f = [x(j),1-x(j)];
        [SE_DT(i,j),SE_DNT(i,j)] = power_slope(n,C,tau,sigma,pi,f,1,2);
    end;
end;

obj   = SE_DT + SE_DNT; % objective function to maximize
[i,j] = index_min(obj); % index of psi that minimizes objective
x(i)         % optimal saturation for lower saturation
1 - x(i)     % optimal saturation for higher saturation
x(j)         % optimal fraction of clusters in lower saturation
1 - x(j)     % optimal fraction of clusters in higher saturation
SE_DT(i,j)  % MDSE_T at optimal 
SE_DNT(i,j)  % MDSE_S at optimal

diary off
diary 'Table2_4.txt'
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                Table 2, Column 4                %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Joint maximization of pi and f across pooled and slope.

clear; clc;
C     = 100;  % number of clusters
n     = 10;   % cluster size
tau   = 0.1;  % variance of cluster error (tau^2)
sigma = 0.9;  % variance of individual error (sigma^2)
x     = [0.01:.01:.99]; % values for optimization
obj_star=Inf; SE_ITT_opt=Inf; SE_SNT_opt=Inf; SE_DT_opt=Inf; SE_DNT_opt=Inf;

for k = 1:length(x)-1;
    for m = k + 1:length(x);
        pi = [0,x(k),x(m)];     
        SE_ITT = ones(length(x),length(x));
        SE_SNT = SE_ITT; SE_DT = SE_ITT; SE_DNT = SE_ITT;

        for i = 1:length(x);
            l = length(x)-i;
            for j = 1:l;
                f = [x(i),x(j),1-x(i)-x(j)]; 
                [SE_ITT(i,j), SE_SNT(i,j),z] = power_pooled(n,C,tau,sigma,pi,f);
                [SE_DT(i,j),z] = power_slope(n,C,tau,sigma,pi,f,2,3);
                [z,SE_DNT(i,j)] = power_slope(n,C,tau,sigma,pi,f,1,3);
            end;
        end;
        
        obj   = SE_ITT + SE_SNT + SE_DT + SE_DNT; % objective function
        [i,j] = index_min(obj); % index of x that minimizes obj
        temp  = SE_ITT(i,j) + SE_SNT(i,j) + SE_DT(i,j) + SE_DNT(i,j);

        if obj_star > temp;
            obj_star = temp;
            psi = x(i); % optimal f(0)
            f_1 = x(j); % optimal f(pi_1)
            pi1 = x(k); % optimal pi_1
            pi2 = x(m); % optimal pi_2
            SE_ITT_opt  = SE_ITT(i,j);  % SE_ITT at optimal values
            SE_SNT_opt  = SE_SNT(i,j);  % SE_SNT at optimal values
            SE_DT_opt = SE_DT(i,j); % SE_DT at optimal values
            SE_DNT_opt = SE_DNT(i,j); % SE_DNT at optimal values
        end;
    end;
end;

[pi1, pi2, psi, f_1, 1-f_1-psi, SE_ITT_opt, SE_SNT_opt, SE_DT_opt, SE_DNT_opt, obj_star]
%Returns SE_DT between two interior saturations 
%Returns SE_DNT between saturations 0 and highest interior saturation

diary off
diary 'Table2_56.txt'
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                Table 2, Column 5                %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% SEs resulting from study design of Banerjee et al.

clear; clc;
n     = 10;        % cluster size
C     = 100;       % number of clusters
tau   = 0.1;       % variance of cluster error (tau^2)
sigma = 0.9;       % variance of individual error (sigma^2)
pi    = [0,0.25,0.5,0.75,1];   % cluster saturations
f     = [0.2,0.2,0.2,0.2,0.2]; % fraction of clusters in each saturation
SE_ITT = 0;SE_SNT = 0;SE_DT = 0;SE_DNT = 0;

[SE_ITT,SE_SNT,z] = power_pooled(n,C,tau,sigma,pi,f);
[SE_DT,z]      = power_slope (n,C,tau,sigma,pi,f,2,5); 
[z,SE_DNT]      = power_slope (n,C,tau,sigma,pi,f,1,4); 
[SE_ITT,SE_SNT,SE_DT,SE_DNT]
%Returns SE_DT between saturations 0.25 and 1
%Returns SE_DNT between saturations 0 and .75

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                Table 2, Column 6                %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% SEs resulting from study design of Sinclair et al.

clear; clc;
n     = 10;        % cluster size
C     = 100;       % number of clusters
tau   = 0.1;       % variance of cluster error (tau^2) 
sigma = 0.9;       % variance of individual error (sigma^2)
pi    = [0,0.1,0.5,1];         % cluster saturations
f     = [0.25,0.25,0.25,0.25]; % fraction of clusters in each saturation
SE_ITT = 0; SE_SNT = 0; SE_DT = 0; SE_DNT = 0;

[SE_ITT,SE_SNT,z] = power_pooled(n,C,tau,sigma,pi,f);
[SE_DT,z]      = power_slope (n,C,tau,sigma,pi,f,2,4);
[z,SE_DNT]      = power_slope (n,C,tau,sigma,pi,f,1,3);
[SE_ITT,SE_SNT,SE_DT,SE_DNT]
%Returns SE_DT between saturations 0.1 and 1
%Returns SE_DNT between saturations 0 and .5

diary off
diary 'Table2_78.txt'
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                Table 2, Column 7-8              %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
clear; clc;
n     = 10;        % cluster size
C     = 100;       % number of clusters
tau   = 0.1;       % variance of cluster error (tau^2)
sigma = 0.9;       % variance of individual error (sigma^2)
pi    = [0,0.33,0.67,1];       % cluster saturations
f     = [0.55,0.15,0.15,0.15]; % fraction of clusters in each saturation

SE_ITT = 0; SE_SNT = 0; SE_DT = 0; SE_DNT = 0;

%%%% Table 2, Column 7
% SEs resulting from study design of Baird et al.
[SE_ITT,SE_SNT,z] = power_pooled(n,C,tau,sigma,pi,f);
[SE_DT,z]      = power_slope (n,C,tau,sigma,pi,f,2,4);
[z,SE_DNT]      = power_slope (n,C,tau,sigma,pi,f,1,3);
[SE_ITT,SE_SNT,SE_DT,SE_DNT]
%Returns SE_DT between saturations 0.33 and 1
%Returns SE_DNT between saturations 0 and .67

%%%% Table 2, Column 8
% Design to minimize SE_SNT subject to:
% 1.)  SE_ITT held no larger than that of Baird et al. (SE<=.095; MDE_T <= 0.2679)
% 2.)  Saturation bins pi held the same as Baird et al.
% 3.)  The share of clusters in interior saturations held equal: f(1/3)=f(2/3)

x     = [0.01:.01:.99]; % values for optimization
SE_ITT = ones(length(x),length(x));SE_SNT = SE_ITT; SE_DT = SE_ITT; SE_DNT = SE_ITT;

for i = 1:length(x);
    l = length(x)-i;
    for j = 1:l;
        f = [x(i),(1-x(i)-x(j))/2,(1-x(i)-x(j))/2, x(j)]; 
        [SE_ITT(i,j), SE_SNT(i,j),z] = power_pooled(n,C,tau,sigma,pi,f);
        [SE_DT(i,j),z] = power_slope (n,C,tau,sigma,pi,f,2,4);
        [z,SE_DNT(i,j)] = power_slope (n,C,tau,sigma,pi,f,1,3);
    end;
end;

SE_SNT_aug  = SE_SNT;
for i = 1:length(x);
    for j = 1:length(x);
        if SE_ITT(i,j) > 0.095;
            SE_SNT_aug(i,j) = Inf;
        end;
    end;
end;
obj   = SE_SNT_aug;      % objective function to maximize
[i,j] = index_min(obj); % index of x that minimizes MDE_S
x(i)            % optimal fraction of clusters in pure control
(1-x(i)-x(j))/2 % optimal fraction of clusters in two middle saturations
x(j)            % optimal fraction of clusters in pure treatment
SE_ITT(i,j)      % MDE_T  at optimal control size
SE_SNT(i,j)      % MDE_S  at optimal control size
SE_DT(i,j)     % MDSE_T at optimal control size
SE_DNT(i,j)     % MDSE_S at optimal control size


